"""
趋势图和热力图数据处理服务
实现数据聚合、预处理、可视化数据生成等功能
"""
import numpy as np
import pandas as pd
from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime, timedelta
import logging
import json

from .models import HeatmapDataIn, HeatmapDataOut, TrendDataIn, TrendDataOut

logger = logging.getLogger(__name__)


class VisualizationService:
    """可视化数据服务"""
    
    def __init__(self):
        self.color_schemes = {
            'viridis': self._viridis_colors,
            'jet': self._jet_colors,
            'gray': self._gray_colors,
            'hot': self._hot_colors
        }
    
    def generate_trend_data(self, signal_data: List[Dict[str, Any]], 
                           metric_name: str = 'signal_strength',
                           aggregation_type: str = 'raw') -> TrendDataOut:
        """生成趋势图数据"""
        try:
            if not signal_data:
                return self._create_empty_trend_data(metric_name, aggregation_type)
            
            # 提取时间序列数据
            time_series = self._extract_time_series(signal_data, metric_name)
            
            # 数据聚合
            if aggregation_type != 'raw':
                time_series = self._aggregate_time_series(time_series, aggregation_type)
            
            # 创建趋势数据对象
            trend_data = TrendDataOut(
                trend_id=f"trend_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}",
                record_id=signal_data[0].get('record_id', 'unknown'),
                device_id=signal_data[0].get('device_id', 'unknown'),
                metric_name=metric_name,
                time_series=time_series,
                aggregation_type=aggregation_type,
                created_at=datetime.utcnow()
            )
            
            logger.info(f"趋势数据生成成功: {metric_name}, 数据点: {len(time_series)}")
            return trend_data
            
        except Exception as e:
            logger.error(f"趋势数据生成失败: {str(e)}")
            return self._create_empty_trend_data(metric_name, aggregation_type)
    
    def generate_heatmap_data(self, signal_data: List[Dict[str, Any]], 
                             data_type: str = 'signal_strength',
                             width: int = 100, height: int = 100,
                             color_scheme: str = 'viridis') -> HeatmapDataOut:
        """生成热力图数据"""
        try:
            if not signal_data:
                return self._create_empty_heatmap_data(data_type, width, height, color_scheme)
            
            # 根据数据类型生成热力图数据
            if data_type == 'signal_strength':
                heatmap_data = self._generate_signal_strength_heatmap(signal_data, width, height)
            elif data_type == 'frequency':
                heatmap_data = self._generate_frequency_heatmap(signal_data, width, height)
            elif data_type == 'time_frequency':
                heatmap_data = self._generate_time_frequency_heatmap(signal_data, width, height)
            else:
                raise ValueError(f"不支持的热力图类型: {data_type}")
            
            # 计算数据范围
            min_value = min(min(row) for row in heatmap_data)
            max_value = max(max(row) for row in heatmap_data)
            
            # 应用颜色方案
            colored_data = self._apply_color_scheme(heatmap_data, color_scheme, min_value, max_value)
            
            # 创建热力图数据对象
            heatmap = HeatmapDataOut(
                heatmap_id=f"heatmap_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}",
                record_id=signal_data[0].get('record_id', 'unknown'),
                device_id=signal_data[0].get('device_id', 'unknown'),
                data_type=data_type,
                width=width,
                height=height,
                data_points=colored_data,
                min_value=min_value,
                max_value=max_value,
                color_scheme=color_scheme,
                created_at=datetime.utcnow()
            )
            
            logger.info(f"热力图数据生成成功: {data_type}, 尺寸: {width}x{height}")
            return heatmap
            
        except Exception as e:
            logger.error(f"热力图数据生成失败: {str(e)}")
            return self._create_empty_heatmap_data(data_type, width, height, color_scheme)
    
    def _extract_time_series(self, signal_data: List[Dict[str, Any]], metric_name: str) -> List[Dict[str, Any]]:
        """提取时间序列数据"""
        time_series = []
        
        for record in signal_data:
            # 获取时间戳
            timestamp = record.get('collected_at')
            if isinstance(timestamp, str):
                timestamp = datetime.fromisoformat(timestamp.replace('Z', '+00:00'))
            elif not isinstance(timestamp, datetime):
                continue
            
            # 获取指标值
            value = record.get(metric_name, 0)
            if value is None:
                continue
            
            time_series.append({
                'timestamp': timestamp.isoformat(),
                'value': float(value)
            })
        
        # 按时间排序
        time_series.sort(key=lambda x: x['timestamp'])
        
        return time_series
    
    def _aggregate_time_series(self, time_series: List[Dict[str, Any]], 
                              aggregation_type: str) -> List[Dict[str, Any]]:
        """聚合时间序列数据"""
        if not time_series:
            return []
        
        # 转换为DataFrame
        df = pd.DataFrame(time_series)
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        df.set_index('timestamp', inplace=True)
        
        # 根据聚合类型进行聚合
        if aggregation_type == 'minute':
            aggregated = df.resample('1T').agg({
                'value': ['mean', 'min', 'max', 'std', 'count']
            }).dropna()
        elif aggregation_type == 'hour':
            aggregated = df.resample('1H').agg({
                'value': ['mean', 'min', 'max', 'std', 'count']
            }).dropna()
        elif aggregation_type == 'day':
            aggregated = df.resample('1D').agg({
                'value': ['mean', 'min', 'max', 'std', 'count']
            }).dropna()
        else:
            return time_series
        
        # 转换回列表格式
        result = []
        for timestamp, row in aggregated.iterrows():
            result.append({
                'timestamp': timestamp.isoformat(),
                'value': float(row[('value', 'mean')]),
                'min': float(row[('value', 'min')]),
                'max': float(row[('value', 'max')]),
                'std': float(row[('value', 'std')]),
                'count': int(row[('value', 'count')])
            })
        
        return result
    
    def _generate_signal_strength_heatmap(self, signal_data: List[Dict[str, Any]], 
                                        width: int, height: int) -> List[List[float]]:
        """生成信号强度热力图"""
        # 提取位置和信号强度数据
        positions = []
        values = []
        
        for record in signal_data:
            lat = record.get('latitude')
            lng = record.get('longitude')
            signal = record.get('signal_strength', -120)
            
            if lat is not None and lng is not None:
                positions.append([float(lat), float(lng)])
                values.append(float(signal))
        
        if not positions:
            return [[-120.0] * width for _ in range(height)]
        
        # 计算边界
        positions = np.array(positions)
        values = np.array(values)
        
        lat_min, lat_max = positions[:, 0].min(), positions[:, 0].max()
        lng_min, lng_max = positions[:, 1].min(), positions[:, 1].max()
        
        # 创建网格
        lat_grid = np.linspace(lat_min, lat_max, height)
        lng_grid = np.linspace(lng_min, lng_max, width)
        
        # 插值生成热力图数据
        heatmap = np.zeros((height, width))
        
        for i, lat in enumerate(lat_grid):
            for j, lng in enumerate(lng_grid):
                # 计算到所有数据点的距离
                distances = np.sqrt((positions[:, 0] - lat)**2 + (positions[:, 1] - lng)**2)
                
                # 使用反距离加权插值
                if np.min(distances) < 1e-6:  # 非常接近数据点
                    heatmap[i, j] = values[np.argmin(distances)]
                else:
                    weights = 1.0 / (distances + 1e-6)
                    heatmap[i, j] = np.average(values, weights=weights)
        
        return heatmap.tolist()
    
    def _generate_frequency_heatmap(self, signal_data: List[Dict[str, Any]], 
                                  width: int, height: int) -> List[List[float]]:
        """生成频率热力图"""
        # 模拟频率数据（实际应该从频谱分析中获取）
        heatmap = np.zeros((height, width))
        
        # 生成模拟的频谱数据
        for i in range(height):
            for j in range(width):
                # 模拟信号强度随频率的变化
                freq_ratio = j / width
                time_ratio = i / height
                
                # 模拟多峰频谱
                value = -120 + 50 * np.exp(-((freq_ratio - 0.3)**2 + (time_ratio - 0.5)**2) * 10)
                value += 30 * np.exp(-((freq_ratio - 0.7)**2 + (time_ratio - 0.3)**2) * 15)
                value += 20 * np.exp(-((freq_ratio - 0.5)**2 + (time_ratio - 0.8)**2) * 8)
                
                heatmap[i, j] = max(-120, min(0, value))
        
        return heatmap.tolist()
    
    def _generate_time_frequency_heatmap(self, signal_data: List[Dict[str, Any]], 
                                       width: int, height: int) -> List[List[float]]:
        """生成时频热力图"""
        if not signal_data:
            return [[-120.0] * width for _ in range(height)]
        
        # 按时间排序
        sorted_data = sorted(signal_data, key=lambda x: x.get('collected_at', ''))
        
        # 创建时频矩阵
        heatmap = np.zeros((height, width))
        
        # 时间轴（高度）
        time_points = height
        # 频率轴（宽度）
        freq_points = width
        
        # 模拟时频分析结果
        for i in range(time_points):
            time_ratio = i / time_points
            
            # 找到对应时间的数据点
            data_index = int(time_ratio * len(sorted_data))
            if data_index >= len(sorted_data):
                data_index = len(sorted_data) - 1
            
            signal_strength = sorted_data[data_index].get('signal_strength', -120)
            
            # 生成该时间点的频谱
            for j in range(freq_points):
                freq_ratio = j / freq_points
                
                # 模拟频谱特征
                base_strength = signal_strength
                freq_modulation = 10 * np.sin(2 * np.pi * freq_ratio * 3)  # 3个主要频率分量
                noise = np.random.normal(0, 2)  # 噪声
                
                heatmap[i, j] = base_strength + freq_modulation + noise
                heatmap[i, j] = max(-120, min(0, heatmap[i, j]))
        
        return heatmap.tolist()
    
    def _apply_color_scheme(self, data: List[List[float]], 
                          color_scheme: str, min_val: float, max_val: float) -> List[List[float]]:
        """应用颜色方案"""
        if color_scheme not in self.color_schemes:
            color_scheme = 'viridis'
        
        # 归一化数据到0-1范围
        data_array = np.array(data)
        if max_val > min_val:
            normalized = (data_array - min_val) / (max_val - min_val)
        else:
            normalized = np.zeros_like(data_array)
        
        # 应用颜色方案
        colored_data = self.color_schemes[color_scheme](normalized)
        
        return colored_data.tolist()
    
    def _viridis_colors(self, normalized_data: np.ndarray) -> np.ndarray:
        """Viridis颜色方案"""
        # 简化的viridis颜色映射
        colors = np.zeros((*normalized_data.shape, 3))
        
        # R通道
        colors[:, :, 0] = 0.267 + 0.004 * normalized_data
        colors[:, :, 0] = np.where(normalized_data > 0.4, 
                                  0.993 - 0.002 * (normalized_data - 0.4) / 0.6, 
                                  colors[:, :, 0])
        
        # G通道
        colors[:, :, 1] = 0.004 + 0.859 * normalized_data
        colors[:, :, 1] = np.where(normalized_data > 0.4, 
                                  0.999 - 0.001 * (normalized_data - 0.4) / 0.6, 
                                  colors[:, :, 1])
        
        # B通道
        colors[:, :, 2] = 0.329 + 0.671 * normalized_data
        
        return colors
    
    def _jet_colors(self, normalized_data: np.ndarray) -> np.ndarray:
        """Jet颜色方案"""
        colors = np.zeros((*normalized_data.shape, 3))
        
        # 简化的jet颜色映射
        colors[:, :, 0] = np.where(normalized_data < 0.25, 
                                  0, 
                                  np.where(normalized_data < 0.5, 
                                          4 * (normalized_data - 0.25), 
                                          np.where(normalized_data < 0.75, 
                                                  1, 
                                                  4 * (1 - normalized_data))))
        
        colors[:, :, 1] = np.where(normalized_data < 0.25, 
                                  4 * normalized_data, 
                                  np.where(normalized_data < 0.5, 
                                          1, 
                                          np.where(normalized_data < 0.75, 
                                                  4 * (0.75 - normalized_data), 
                                                  0)))
        
        colors[:, :, 2] = np.where(normalized_data < 0.5, 
                                  0, 
                                  np.where(normalized_data < 0.75, 
                                          4 * (normalized_data - 0.5), 
                                          1))
        
        return colors
    
    def _gray_colors(self, normalized_data: np.ndarray) -> np.ndarray:
        """灰度颜色方案"""
        colors = np.zeros((*normalized_data.shape, 3))
        colors[:, :, 0] = normalized_data  # R
        colors[:, :, 1] = normalized_data  # G
        colors[:, :, 2] = normalized_data  # B
        return colors
    
    def _hot_colors(self, normalized_data: np.ndarray) -> np.ndarray:
        """Hot颜色方案"""
        colors = np.zeros((*normalized_data.shape, 3))
        
        # R通道：从0到1
        colors[:, :, 0] = np.clip(3 * normalized_data, 0, 1)
        
        # G通道：从0到1
        colors[:, :, 1] = np.clip(3 * normalized_data - 1, 0, 1)
        
        # B通道：从0到1
        colors[:, :, 2] = np.clip(3 * normalized_data - 2, 0, 1)
        
        return colors
    
    def _create_empty_trend_data(self, metric_name: str, aggregation_type: str) -> TrendDataOut:
        """创建空的趋势数据"""
        return TrendDataOut(
            trend_id=f"empty_trend_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}",
            record_id="unknown",
            device_id="unknown",
            metric_name=metric_name,
            time_series=[],
            aggregation_type=aggregation_type,
            created_at=datetime.utcnow()
        )
    
    def _create_empty_heatmap_data(self, data_type: str, width: int, height: int, 
                                 color_scheme: str) -> HeatmapDataOut:
        """创建空的热力图数据"""
        empty_data = [[-120.0] * width for _ in range(height)]
        
        return HeatmapDataOut(
            heatmap_id=f"empty_heatmap_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}",
            record_id="unknown",
            device_id="unknown",
            data_type=data_type,
            width=width,
            height=height,
            data_points=empty_data,
            min_value=-120.0,
            max_value=-120.0,
            color_scheme=color_scheme,
            created_at=datetime.utcnow()
        )